package com.feidee.fd.sml.algorithm.feature

import com.feidee.fd.sml.algorithm.component.feature.{ElementwiseProductor, ElementwiseProductorParam}
import com.feidee.fd.sml.algorithm.util.{TestingDataGenerator, ToolClass}
import org.apache.spark.ml.linalg.Vectors
import org.scalatest.FunSuite

/**
  * @author YongChen
  * @date 2018/12/25 10:18
  * @description
  * @email yong_chen@sui.com
  */
class ElementwiseProductSuite extends FunSuite {

  val paramStr: String =
        """
     |{
     |	"input_pt": "a",
     |	"output_pt": "b",
     | "hive_table": "",
     | "flow_time": "",
     |	"inputCol": "vector",
     |	"outputCol": "scaledFeatures",
     | "preserveCols": "",
     |	"modelPath": "",
     |  "scalingVec" :[0.0, 1.0, 2.0]
     | }
    """.stripMargin

  val ep = new ElementwiseProductor()
  val param: ElementwiseProductorParam = ep.parseParam(new ToolClass().encrypt(paramStr))

  test("ElementwiseProduct parameter construction test") {
    assert(
          "a".equals(param.input_pt)
        & "b".equals(param.output_pt)
        & "".equals(param.hive_table)
        & "".equals(param.flow_time)
        & "vector".equals(param.inputCol)
        & "scaledFeatures".equals(param.outputCol)
        & "".equals(param.preserveCols)
        & "".equals(param.modelPath)
    )
  }

  test("ElementwiseProduct training function test") {
    val dataFrame = TestingDataGenerator.spark.createDataFrame(Seq(
      ("a", Vectors.dense(1.0, 2.0, 3.0)),
      ("b", Vectors.dense(4.0, 5.0, 6.0)))).toDF("id", "vector")

    val expectRes = TestingDataGenerator.spark.createDataFrame(Seq(
      ("a", Vectors.dense(0.0, 2.0, 6.0)),
      ("b", Vectors.dense(0.0, 5.0, 12.0)))).toDF("id", param.outputCol)

    val res = ep.train(param, dataFrame).transform(dataFrame)

    assert(res.select(param.outputCol).except(expectRes.select(param.outputCol)).count() == 0)
  }
}
